import os
import requests

# 下载模型实例，国内源在url前添加 https://mirror.ghproxy.com/ 
# 其他链接在 resources 目录下 rembg_model.txt 文件中 
def download_model():
    model_url = "https://mirror.ghproxy.com/https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx"
    save_dir = os.path.expanduser("~/.u2net")
    save_path = os.path.join(save_dir, "u2net.onnx")
    
    # 创建保存目录
    os.makedirs(save_dir, exist_ok=True)
    
    # 下载模型
    print(f"正在下载模型到: {save_path}")
    response = requests.get(model_url, stream=True)
    total_size = int(response.headers.get('content-length', 0))
    
    with open(save_path, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            if chunk:
                f.write(chunk)
    
    print("模型下载完成！")

if __name__ == "__main__":
    download_model() 